/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include <sstream>
#include "../../common/data_utils.h"
#include "kernel_tiling/kernel_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#include "op_host/matmul_nbuffer33_custom_tiling.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"

extern void matmul_nbuffer33_custom_do(uint32_t coreDim, void* stream,
                             uint8_t *a, uint8_t *b, uint8_t *bias, uint8_t *c,
                             uint8_t *workspace, uint8_t *tiling);
#else
#include "tikicpulib.h"
extern "C" void matmul_nbuffer33_custom(uint8_t *a, uint8_t *b, uint8_t *bias, uint8_t *c,
                              uint8_t *workspace, uint8_t *tiling);
#endif

namespace {
bool IS_BIAS = false;
bool IS_A_TRANS = false;
bool IS_B_TRANS = false;
}

namespace MatmulHost {
static size_t GetSysWorkSpaceSize()
{
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
    return static_cast<size_t>(ascendcPlatform->GetLibApiWorkSpaceSize());
}

// CPU debug mode
#ifdef ASCENDC_CPU_DEBUG
void TestMatmulCpu(const MatmulHost::MatmulCaseParams &testCaseParams)
{
    int64_t m = testCaseParams.m;
    int64_t n = testCaseParams.n;
    int64_t k = testCaseParams.k;
    size_t x1FileSize = m * k * sizeof(uint16_t);  // uint16_t represent half
    size_t x2FileSize = k * n * sizeof(uint16_t);  // uint16_t represent half
    size_t yFileSize = m * n * sizeof(float);
    size_t biasFileSize = 1 * n * sizeof(float);
    // Query workspace size
    size_t workspaceSize = GetSysWorkSpaceSize();
    uint8_t* x1 = (uint8_t*)AscendC::GmAlloc(x1FileSize);
    uint8_t* x2 = (uint8_t*)AscendC::GmAlloc(x2FileSize);
    uint8_t* bias; // IS_BIAS = false
    uint8_t* y = (uint8_t*)AscendC::GmAlloc(yFileSize);
    uint8_t* workspace = (uint8_t*)AscendC::GmAlloc(workspaceSize);
    ReadFile("../input/x1_gm.bin", x1FileSize, x1, x1FileSize);
    ReadFile("../input/x2_gm.bin", x2FileSize, x2, x2FileSize);
    if (IS_BIAS) {
        bias = (uint8_t*)AscendC::GmAlloc(biasFileSize);
        ReadFile("../input/bias_gm.bin", biasFileSize, bias, biasFileSize);
    }
    size_t tilingFileSize = sizeof(TCubeTiling);
    uint8_t* tiling = (uint8_t*)AscendC::GmAlloc(tilingFileSize);
    // Calculate Tiling
    const auto& tilingData = MatmulHost::GenerateTiling(testCaseParams);
    memcpy_s(tiling, tilingFileSize, &tilingData, tilingFileSize);
    ICPU_RUN_KF(matmul_nbuffer33_custom, tilingData.usedCoreNum, x1, x2, bias, y, workspace, tiling);
    WriteFile("../output/output.bin", y, yFileSize);
    AscendC::GmFree((void*)x1);
    AscendC::GmFree((void*)x2);
    AscendC::GmFree((void*)y);
    AscendC::GmFree((void*)workspace);
    AscendC::GmFree((void*)tiling);
    if (IS_BIAS) {
        AscendC::GmFree((void*)bias);
    }
}
// NPU
#else
void MatmulOp(uint8_t* x1, uint8_t* x2, uint8_t* y, uint8_t* bias, MatmulHost::MatmulCaseParams testCaseParams,
    void* stream = nullptr)
{
    // Init args
    uint8_t *workspaceDevice;

    // Query workspace size
    size_t workspaceSize = GetSysWorkSpaceSize();

    // Allocate workspace on device
    CHECK_ACL(aclrtMalloc((void **)&workspaceDevice, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));

    // Initialize kernel with arguments and workspace pointer
    uint8_t* tilingHost;
    uint8_t* tilingDevice;
    size_t tilingFileSize = sizeof(TCubeTiling);
    // Calculate Tiling
    const auto tilingData = MatmulHost::GenerateTiling(testCaseParams);
    CHECK_ACL(aclrtMallocHost((void **)(&tilingHost), tilingFileSize));
    CHECK_ACL(aclrtMalloc((void **)&tilingDevice, tilingFileSize,
                          ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemcpy(tilingHost, tilingFileSize, &tilingData,
                          tilingFileSize, ACL_MEMCPY_HOST_TO_HOST));
    CHECK_ACL(aclrtMemcpy(tilingDevice, tilingFileSize, tilingHost,
                          tilingFileSize, ACL_MEMCPY_HOST_TO_DEVICE));
    // Launch kernel
    matmul_nbuffer33_custom_do(tilingData.usedCoreNum, stream, x1, x2, bias, y, workspaceDevice, tilingDevice);
    CHECK_ACL(aclrtFreeHost(tilingHost));
    CHECK_ACL(aclrtFree(workspaceDevice));
    CHECK_ACL(aclrtFree(tilingDevice));
}

void TestAclInit(aclrtContext &context, aclrtStream &stream, int64_t &deviceId)
{
    CHECK_ACL(aclInit(nullptr));
    CHECK_ACL(aclrtSetDevice(deviceId));
    CHECK_ACL(aclrtCreateContext(&context, deviceId));
    CHECK_ACL(aclrtCreateStream(&stream));
}

void TestAclDeInit(aclrtContext &context, aclrtStream &stream, int64_t &deviceId)
{
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtDestroyContext(context));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
}

void TestMatmul(const MatmulHost::MatmulCaseParams &testCaseParams)
{
    size_t x1FileSize = static_cast<size_t>(testCaseParams.m * testCaseParams.k) * sizeof(uint16_t);  // uint16_t represent half
    size_t x2FileSize = static_cast<size_t>(testCaseParams.k * testCaseParams.n) * sizeof(uint16_t);  // uint16_t represent half
    size_t yFileSize = static_cast<size_t>(testCaseParams.m * testCaseParams.n) * sizeof(float);
    size_t biasFileSize = static_cast<size_t>(1 * testCaseParams.n) * sizeof(float);

    aclrtContext context;
    aclrtStream stream = nullptr;
    int64_t deviceId = 0;
    TestAclInit(context, stream, deviceId);

    uint8_t *x1Host;
    uint8_t *x1Device;
    CHECK_ACL(aclrtMallocHost((void **)(&x1Host), x1FileSize));
    CHECK_ACL(aclrtMalloc((void **)&x1Device, x1FileSize, ACL_MEM_MALLOC_HUGE_FIRST));
    ReadFile("../input/x1_gm.bin", x1FileSize, x1Host, x1FileSize);
    CHECK_ACL(aclrtMemcpy(x1Device, x1FileSize, x1Host, x1FileSize, ACL_MEMCPY_HOST_TO_DEVICE));

    uint8_t *x2Host;
    uint8_t *x2Device;
    CHECK_ACL(aclrtMallocHost((void **)(&x2Host), x2FileSize));
    CHECK_ACL(aclrtMalloc((void **)&x2Device, x2FileSize, ACL_MEM_MALLOC_HUGE_FIRST));
    ReadFile("../input/x2_gm.bin", x2FileSize, x2Host, x2FileSize);
    CHECK_ACL(aclrtMemcpy(x2Device, x2FileSize, x2Host, x2FileSize, ACL_MEMCPY_HOST_TO_DEVICE));

    uint8_t *biasHost = nullptr;
    uint8_t *biasDevice = nullptr;
    if (IS_BIAS) {
        CHECK_ACL(aclrtMallocHost((void **)(&biasHost), biasFileSize));
        CHECK_ACL(aclrtMalloc((void **)&biasDevice, biasFileSize, ACL_MEM_MALLOC_HUGE_FIRST));
        ReadFile("../input/bias_gm.bin", biasFileSize, biasHost, biasFileSize);
        CHECK_ACL(aclrtMemcpy(biasDevice, biasFileSize, biasHost, biasFileSize, ACL_MEMCPY_HOST_TO_DEVICE));
    }
    uint8_t *yHost = nullptr;
    uint8_t *yDevice = nullptr;
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), yFileSize));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, yFileSize, ACL_MEM_MALLOC_HUGE_FIRST));

    MatmulOp(x1Device, x2Device, yDevice, biasDevice, testCaseParams, stream);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(yHost, yFileSize, yDevice, yFileSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("../output/output.bin", yHost, yFileSize);

    if (IS_BIAS) {
        CHECK_ACL(aclrtFree(biasDevice));
        CHECK_ACL(aclrtFreeHost(biasHost));
    }
    CHECK_ACL(aclrtFree(x1Device));
    CHECK_ACL(aclrtFreeHost(x1Host));
    CHECK_ACL(aclrtFree(x2Device));
    CHECK_ACL(aclrtFreeHost(x2Host));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFreeHost(yHost));
    TestAclDeInit(context, stream, deviceId);
}
#endif
} // namespace MatmulHost

int32_t main(int32_t argc, const char *args[])
{
    int64_t problem[12] = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
    for (int32_t i = 1; i < argc && i < 13; ++i) { // 13
        std::stringstream ss(args[i]);
        ss >> problem[i - 1];
    }
    const char* socVersion = SOC_VERSION;
    auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(socVersion);
    MatmulHost::MatmulCaseParams testCaseParams{static_cast<int32_t>(ascendcPlatform->GetCoreNumAic()),
        static_cast<int32_t>(problem[0]), static_cast<int32_t>(problem[1]), static_cast<int32_t>(problem[2]),
        IS_BIAS, IS_A_TRANS, IS_B_TRANS,
        static_cast<int32_t>(problem[3]), static_cast<int32_t>(problem[4]), static_cast<int32_t>(problem[5]),
        static_cast<int32_t>(problem[6]), static_cast<int32_t>(problem[7]), static_cast<int32_t>(problem[8]),
        static_cast<int32_t>(problem[9]), static_cast<int32_t>(problem[10]), static_cast<int32_t>(problem[11])};
#ifdef ASCENDC_CPU_DEBUG
    MatmulHost::TestMatmulCpu(testCaseParams);
#else
    MatmulHost::TestMatmul(testCaseParams);
#endif
    return 0;
}